#!/bin/bash
#
# Drake script to set up the Python venv based on information in setup/python.
# Uses https://pdm-project.org/ under the hood.
#
# Users must NOT run this manually.

set -eux -o pipefail

# Globals (will be set from command line arguments).
python=
repository=
symlink=

# Create virtual environment at specified path (if it doesn't exist).
mkvenv() {
    if [ ! -d "$1" ]; then
        "${python}" -m venv "$1"
    fi
}

# Process command line arguments.
while [ "${1:-}" != "" ]; do
    case "$1" in
        --python)
            # The python interpreter to use.
            readonly python="$2"
            shift
            ;;
        --repository)
            # The bazel repository rule root to interface with. We'll use the
            # response file found there (i.e. the file that supplies additional
            # arguments to be given to `pdm sync`), and we'll communicate our
            # result by writing back a `venv_python3.txt` file to the same
            # location, containing the path to our venv.drake/bin/python3.
            readonly repository="$2"
            shift
            ;;
        --symlink)
            # Create a symlink to the virtual environment in the Drake source
            # tree. This is intended to be used by wheel builds, which may,
            # outside of the Bazel build, need access to tools installed in the
            # virtual environment.
            readonly symlink="venv"
            ;;
        *)
            echo 'Invalid command line argument' >&2
            exit 5
    esac
    shift
done
if [[ -z "${python}" ]]; then
    echo "error: --python is required"
    exit 1
fi
if [[ -z "${repository}" ]]; then
    echo "error: --repository is required"
    exit 1
fi

# Place the venv(s) in a sibling directory to the output base. That should be a
# suitable disk location for build artifacts, but without polluting the actual
# output base that Bazel owns.
readonly drake_root="$(cd "$(dirname $0)/../../.." && pwd)"
readonly bazel_output_base="$(cd "${repository}/../.." && pwd)"
readonly drake_python="${bazel_output_base}.drake_python"
mkdir -p "${drake_python}"

# Install PDM into a virtual environment. We segregate PDM from the environment
# it is managing, so that changes to the managed environment cannot break PDM.
readonly setup="${drake_root}/setup/python"
readonly venv_pdm="${drake_python}/venv.pdm"
mkvenv "${venv_pdm}"
"${venv_pdm}/bin/pip" install -U -r "${setup}/requirements.txt"

# Don't nag about new versions of PDM; we'll update the above pin when we want
# to use something newer, and otherwise want to use the pinned version, so
# being informed about newer versions is just noise.
export PDM_CHECK_UPDATE=0

# Prepare the PDM "project directory".
readonly project="${drake_python}/project"
mkdir -p "${project}"
ln -nsf "${setup}/pyproject.toml" "${project}/pyproject.toml"
ln -nsf "${setup}/pdm.lock" "${project}/pdm.lock"

# Don't nag about new versions of PDM; venv_upgrade will check for those.
# Aside from that, we want to use the pinned version, so being informed about
# newer versions is just noise.
export PDM_CHECK_UPDATE=0

# Prepare the venv that will hold Drake's requirements.
readonly venv_drake="${drake_python}/venv.drake"
mkvenv "${venv_drake}"
"${venv_pdm}/bin/pdm" use -p "${project}" -f "${venv_drake}"

# Ask PDM to install whatever dependencies we need.
readonly install_args=( $(cat "${repository}/@pdm-install-args") )
"${venv_pdm}/bin/pdm" sync \
    -p "${project}" \
    --clean-unselected \
    "${install_args[@]}"

# Compute the checksum of our requirements file.
readonly checksum=$(cd "${project}" && "${python}" <<EOF
import hashlib
from pathlib import Path
data = Path("pdm.lock").read_bytes()
digest = hashlib.sha256(data).hexdigest()
print(digest, flush=True)
EOF
)
if [[ -z "${checksum}" ]]; then
    echo "error: computing checksum failed"
    exit 1
fi

# Symlink our venv path for the repository.bzl to use. It's directory name
# inside of our repository is a hash of the requirements file so that Bazel
# will invalidate all Python targets when the checksum changes.
ln -nsf "${venv_drake}" "${repository}/${checksum}"

# Tell repository.bzl which path to use.
echo -n "${checksum}/bin/python3" > "${repository}/venv_python3.txt"

# Symlink our venv path for wheel builds to use (if requested; see
# documentation of --symlink, above).
if [[ -n "${symlink}" ]]; then
    ln -nsf "${venv_drake}" "${drake_root}/${symlink}"
fi
